package com.ruoyi.AI.service;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import jakarta.annotation.Resource;
import java.util.ArrayList;
import java.util.List;

@Slf4j
@Service
@RequiredArgsConstructor
public class AiService {

    @Resource
    private OllamaChatModel ollamaChatModel;


    /**
     * 最大消息记录数
     */
    private final static Integer MAX_SIZE = 10;

    /**
     * 消息记录
     */
    private List<Message> messages = new ArrayList<Message>();

    /**
     * 初始化存入系统消息
     */
    private void addSystemMessage() {
        String message = "";
        Message systemMessage = new SystemMessage(message);
        messages.add(systemMessage);
    }

    /**
     * 存储用户发送的消息
     * @param message
     */
    private void addUserMessage(String message) {
        Message userMessage = new UserMessage(message);
        messages.add(userMessage);
    }

    /**
     * 存储AI回复的消息
     * @param message
     */
    private void addAssistantMessage(String message) {
        Message assistantMessage = new AssistantMessage(message);
        messages.add(assistantMessage);
    }

    /**
     * 聊天接口
     * @param message
     * @return
     */
    public String chat(String message) {
        addUserMessage(message);
        String result = ollamaChatModel.call(new Prompt(message)).getResult().getOutput().getText();
        //aiClient.call(new Prompt(messages)).getResult().getOutput().getText();
        addAssistantMessage(result);
        update();
        return result;
    }

    /**
     * 流式聊天接口
     * @param message
     * @return
     */
    public Flux<String> chatStream(String message) {
        addUserMessage(message);
        StringBuilder fullReply = new StringBuilder();
        Flux<String> fluxResult = ollamaChatModel.stream(new Prompt(messages))
                .flatMap(response -> {
                    String reply = response.getResult().getOutput().getText();
                    fullReply.append(reply);//拼接回复内容
                    return Flux.just(reply);
                }).doOnComplete(() -> {
                    //监听流式响应完成，完整回复存入消息记录
                    System.out.println(fullReply);
                    addAssistantMessage(String.valueOf(fullReply));
                });
        update();
        return fluxResult;
    }

    /**
     * 更新聊天记录
     */
    private void update() {
        if (messages.size() > MAX_SIZE) {
            messages = messages.subList(messages.size() - MAX_SIZE, messages.size());
        }
    }

    public String sendPythonToAi(String message){
        return null;
    }
}
